package com.example.shardingsphere.sharding;

import com.example.shardingsphere.SpringContextUtil;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * describe: 自定义库分片算法
 *
 * @author lianying
 * @date 2023/05/05
 */
public class CustomDbComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm {

    @Override
    public Collection<String> doSharding(Collection collection, ComplexKeysShardingValue complexKeysShardingValue) {
        UserIdHandler userIdHandler= SpringContextUtil.getBean(UserIdHandler.class);
        ArrayList<String> databases=new ArrayList(collection);
        Map<String, Collection<Comparable>> map = complexKeysShardingValue.getColumnNameAndShardingValuesMap();
        Collection<Comparable> id = map.get("id");
        //判断所属db
        Iterator<Comparable> iterator = id.iterator();
        //配置table集合
        Collection<String> dbSet=new HashSet<>();
        while (iterator.hasNext()){
            int index =  userIdHandler.getDbIndex(Long.valueOf(String.valueOf(iterator.next())));
            dbSet.add(databases.get(index));
        }
        return dbSet;
    }

}
